#!/usr/bin/env python
from itertools import product

class TSP:
    def __init__(self, graph):
        """graph = networkx undirected graph object. each edge has
        attribute "length"."""
        self.graph = graph.copy()

    def __len__(self):
        return len(self.graph)

    def tour_length(self, permutation):
        assert(len(permutation) == len(set(permutation)) == len(self.graph))
        acc = self.graph[permutation[0]][permutation[-1]]['length']
        for u, v in zip(permutation, permutation[1:]):
            acc += self.graph[u][v]['length']
        return acc

    def __repr__(self):
        return """<TSP {0} nodes>""".format(len(self.graph))

    def generate_ip(self):
        import gurobipy as grb
        m = grb.Model(self.__repr__())

        # convert to directed
        g = self.graph.to_directed()

        # edge varibles
        for u, v, data in g.edges(data=True):
            data['edge_var'] = m.addVar(vtype=grb.GRB.BINARY,
                name="edge_{0}_{1}".format(u, v), obj=data['length'])
        m.update()

        # add in/out constraints
        for node in g:
            out_cons = grb.quicksum(g[node][v]['edge_var']
                for _, v in g.out_edges(node))
            in_cons = grb.quicksum(g[v][node]['edge_var']
                for v, _ in g.in_edges(node))
            m.addConstr(out_cons == 1, name="out_cons_{0}".format(node))
            m.addConstr(in_cons == 1, name="in_cons_{0}".format(node))
        m.update()

        # add subtour constraints
        first_city = g.nodes()[0]
        for node, data in g.nodes(data=True):
            if node == first_city:
                continue
            data['seq'] = m.addVar(name="seq_{0}".format(node))
        m.update()

        for i, j in product(g.nodes(), g.nodes()):
            if i == first_city or j == first_city or i == j:
                continue
            u_i, u_j = g.node[i]['seq'], g.node[j]['seq']
            n = len(g)
            m.addConstr(u_i - u_j + n * g[i][j]['edge_var'] <= n - 1,
                name="subtour_{0}_{1}".format(i, j))
        m.update()

        return m

    @staticmethod
    def from_file(filename):
        from networkx import Graph
        g = Graph()
        with open(filename) as f:
            read_edges = False
            read_nodes = False
            for line in f:
                if line.startswith("EDGE_LENGTH_SECTION"):
                    read_edges = True
                    read_nodes = False
                    continue
                if line.startswith("NODE_COORD_SECTION"):
                    read_edges = False
                    read_nodes = True
                    continue
                if "EOF" in line:
                    break
                if read_nodes:
                    n, x, y = line.strip().split(" ")
                    g.add_node(int(n) - 1, pos=(float(x), float(y)))
                if read_edges:
                    u, v, l = map(int, line.split())
                    g.add_edge(u - 1, v - 1, length=l)
        return TSP(g)


stuff_to_draw = [(0, ([0, 13, 5, 9, 12, 8, 10, 6, 14, 2, 11, 4, 7, 1, 3, 15], 10872)),
(1, ([2, 3, 1, 12, 0, 11, 15, 7, 14, 4, 13, 6, 5, 8, 10, 9], 10026)),
(3, ([4, 14, 11, 15, 7, 3, 1, 12, 5, 9, 10, 8, 6, 13, 2, 0], 9926)),
(5, ([8, 12, 13, 2, 3, 1, 0, 11, 15, 7, 14, 4, 6, 5, 9, 10], 9434)),
(7, ([8, 4, 14, 11, 15, 7, 3, 12, 13, 2, 1, 0, 6, 5, 10, 9], 9203)),
(8, ([8, 4, 14, 11, 15, 7, 3, 12, 13, 2, 1, 0, 6, 5, 9, 10], 9060)),
(9, ([4, 6, 8, 12, 13, 14, 11, 15, 7, 3, 2, 1, 0, 5, 9, 10], 8699)),
(10, ([8, 12, 13, 11, 6, 3, 1, 2, 0, 15, 7, 14, 4, 5, 9, 10], 8300)),
(16, ([8, 4, 14, 11, 12, 13, 15, 7, 2, 1, 3, 0, 6, 5, 9, 10], 8052)),
(25, ([8, 4, 14, 12, 13, 11, 15, 7, 2, 1, 3, 0, 6, 5, 9, 10], 8008)),
(30, ([8, 5, 4, 14, 12, 13, 11, 15, 7, 2, 1, 3, 0, 6, 9, 10], 7892)),
(38, ([8, 5, 4, 14, 12, 13, 11, 15, 7, 3, 2, 1, 0, 6, 9, 10], 7891)),
(43, ([8, 5, 4, 14, 12, 13, 11, 6, 0, 7, 3, 1, 2, 15, 9, 10], 7750)),
(51, ([8, 5, 4, 14, 12, 13, 6, 11, 0, 7, 3, 1, 2, 15, 9, 10], 7675)),
(62, ([10, 5, 4, 14, 12, 13, 11, 6, 15, 0, 7, 3, 1, 2, 9, 8], 7514)),
(65, ([8, 10, 4, 14, 11, 12, 13, 15, 7, 3, 1, 2, 0, 9, 6, 5], 7374)),
(66, ([8, 10, 4, 14, 11, 12, 13, 5, 6, 0, 7, 3, 1, 2, 15, 9], 7148)),
(69, ([8, 10, 4, 14, 12, 13, 5, 6, 11, 0, 7, 3, 1, 2, 15, 9], 7031)),
(117, ([8, 10, 4, 14, 5, 6, 12, 13, 11, 15, 0, 7, 3, 1, 2, 9], 6959)),
(122, ([8, 10, 4, 14, 5, 6, 12, 13, 11, 0, 7, 3, 1, 2, 15, 9], 6935)),
(127, ([8, 10, 4, 14, 5, 6, 11, 12, 13, 0, 7, 3, 1, 2, 15, 9], 6859))
]


def graph_from_tour(instance, tour):
    import networkx as nx
    g = nx.Graph()
    for node, data in instance.graph.nodes(data=True):
        g.add_node(node, **data)
    for (u, v) in zip(tour, tour[1:]) + [(tour[-1], tour[0])]:
        g.add_edge(u, v)
    return g


def draw_tours(instance):
    import matplotlib.pyplot as plt
    import networkx as nx
    for i, (iteration, (tour, quality)) in enumerate(stuff_to_draw):
        plt.close()
        g = graph_from_tour(instance, tour)
        pos = dict((u, d['pos']) for (u, d) in g.nodes(data=True))
        nx.draw(g, pos=pos, node_size=200, width=3)
        plt.annotate("iter: {0}".format(iteration), xy=(40, -3))
        plt.annotate("len: {0}".format(quality), xy=(40, -4.4))
        plt.savefig("stuff_{1}_{0:02d}.png".format(i, iteration))
